package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo05ImageModelTrain {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo05ImageModelTrain")
      .getOrCreate()

    import spark.implicits._

    // 1、加载数据
    val imageDF: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", "784")
      .load("spark/data/mllib/data/imagelibsvm")

    // 2、将数据分为训练集以及测试集
    val randomSplitArr: Array[DataFrame] = imageDF.randomSplit(Array(0.8, 0.2))
    val trainDF: DataFrame = randomSplitArr(0)
    val testDF: DataFrame = randomSplitArr(1)

    // 3、基于数据选择合适的算法，再确定模型
    val lr: LogisticRegression = new LogisticRegression()
      .setMaxIter(10)
      //      .setRegParam(0.3)
      //      .setElasticNetParam(0.8)
      .setFitIntercept(true) // 设置是否有截距
      .setFamily("multinomial")

    // 4、将训练集带入模型 进行训练
    val lRModel: LogisticRegressionModel = lr.fit(trainDF)

    // 5、将测试集带入模型 进行预测
    val predictionDF: DataFrame = lRModel.transform(testDF)
    predictionDF.show()

    // 6、基于预测的结果 评估模型
    // 准确率: 预测准确地/总的数据条数
    println(s"模型的准确率为：${predictionDF.where($"label" === $"prediction").count().toDouble / predictionDF.count()}")

    // 7、保存模型
    lRModel.write.overwrite().save("spark/data/mllib/image")

  }

}
